/*
 * MCMCMC.java
 *
 * Copyright (C) 2002-2006 Alexei Drummond and Andrew Rambaut
 *
 * This file is part of BEAST.
 * See the NOTICE file distributed with this work for additional
 * information regarding copyright ownership and licensing.
 *
 * BEAST is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 *  BEAST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with BEAST; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301  USA
 */

package dr.inference.mcmcmc;

import dr.inference.loggers.Logger;
import dr.inference.loggers.MCLogger;
import dr.inference.markovchain.MarkovChain;
import dr.inference.markovchain.MarkovChainListener;
import dr.inference.model.Likelihood;
import dr.inference.model.Model;
import dr.inference.operators.*;
import dr.inference.prior.Prior;
import dr.inference.mcmc.MCMC;
import dr.inference.mcmc.MCMCCriterion;
import dr.inference.mcmc.MCMCOptions;
import dr.util.NumberFormatter;
import dr.math.MathUtils;

import java.util.ArrayList;

/**
 * An MCMC analysis that estimates parameters of a probabilistic model.
 *
 * @author Andrew Rambaut
 *
 * @version $Id: MCMCMC.java,v 1.12 2005/01/10 10:56:59 rambaut Exp $
 *
 */
public class MCMCMC implements Runnable {

    public MCMCMC(MCMC[] mcmcs, MCMCMCOptions mcmcmcOptions) {

        this.mcmcmcOptions = mcmcmcOptions;

        if (mcmcmcOptions.getChainTemperatures()[0] != 1.0) {
            throw new RuntimeException("The first chain in the array should be cold (temperature = 1.0)");
        }

	    coldChain = 0;

        this.mcmcOptions = mcmcs[coldChain].getOptions();

        // Get all the loggers out of all the chains. We will only use the
        // loggers of the cold chain but we need to swap the formatters around
        // so that which every chain is cold always writes to the same destination.
        mcLoggers = new MCLogger[mcmcs.length][];
        for (int i = 0; i < mcmcs.length; i++) {
            Logger[] loggers = mcmcs[i].getLoggers();
            mcLoggers[i] = new MCLogger[loggers.length];
            for (int j = 0; j < loggers.length; j++) {
                mcLoggers[i][j] = (MCLogger)loggers[j];
            }
            if (mcLoggers[i] == null) {
                throw new RuntimeException("There are no loggers in the MCMC chains.");
            }
        }

        // Get all the operator schedules. The tuning values of these must be swapped
        // around as the temperatures are swapped.
		schedules = new OperatorSchedule[mcmcs.length];
        for (int i = 0; i < schedules.length; i++) {
            schedules[i] = mcmcs[i].getOperatorSchedule();
        }

        chains = new MarkovChain[mcmcs.length];

        chains[0] = mcmcs[0].getMarkovChain();
        for (int i = 1; i < chains.length; i++) {
            chains[i] = mcmcs[i].getMarkovChain();
	        MCMCCriterion acceptor = ((MCMCCriterion)chains[i].getAcceptor());
	        acceptor.setTemperature(mcmcmcOptions.getChainTemperatures()[i]);
        }

    }

    public void run() {
        currentState = 0;

        timer.start();

        if (isPreBurninNeeded()) {
            int preBurnin = mcmcOptions.getPreBurnin();
            if (preBurnin > 0) {
                MarkovChainListener burninListener = new BurninListener(preBurnin);

                chains[coldChain].addMarkovChainListener(burninListener);
                runChains(preBurnin, true);
                chains[coldChain].removeMarkovChainListener(burninListener);
                resetChains();
            }
        }

        MCLogger[] loggers = mcLoggers[coldChain];
        ArrayList[] logFormatters = new ArrayList[loggers.length];

        for (int i = 0; i < loggers.length; i++) {
            // Start the logging for the cold chain
            loggers[i].startLogging();

            // Now get the formatters (destinations) for the cold chains loggers
            logFormatters[i] = loggers[i].getFormatters();
        }

        int oldColdChain = 0;

        // Set the cold chain's loggers with the formatters (destinations) of
        // the original cold chain
        loggers = mcLoggers[coldChain];
        for (int i = 0; i < loggers.length; i++) {
            loggers[i].setFormatters(logFormatters[i]);
        }

        chains[coldChain].addMarkovChainListener(chainListener);

        MCMCMCRunner[] threads = new MCMCMCRunner[chains.length];
        for (int i = 0; i < chains.length; i++) {
            threads[i] = new MCMCMCRunner(chains[i], mcmcmcOptions.getSwapChainsEvery(), getChainLength(), false);
            threads[i].start();
        }


        while (chains[coldChain].getCurrentLength() < getChainLength()) {

            // wait for all the threads to complete their alloted chain length
            boolean allDone = true;
            do {
                try {
                    Thread.sleep(10);
                } catch (InterruptedException e) { }

                allDone = true;
                for (int i = 0; i < chains.length; i++) {
                    if (!threads[i].isChainDone()) {
                        allDone = false;
                    }
                }
            } while (!allDone);

            if (chains[coldChain].getCurrentLength() < getChainLength()) {
                // attempt to swap two chains' temperatures
                coldChain = swapChainTemperatures();

                // if the cold chain was involved in a swap then we need to change the
                // listener that does the logging and the destinations for the loggers.
                if (coldChain != oldColdChain) {

                    chains[oldColdChain].removeMarkovChainListener(chainListener);

                    // Set the cold chain's loggers with the formatters (destinations) of
                    // the original cold chain
                    loggers = mcLoggers[coldChain];
                    for (int j = 0; j < loggers.length; j++) {
                        loggers[j].setFormatters(logFormatters[j]);
                    }

                    chains[coldChain].addMarkovChainListener(chainListener);

                    oldColdChain = coldChain;
                }

                for (int i = 0; i < chains.length; i++) {
	                threads[i].continueChain();
                }
            }

        }

        finish();

        timer.stop();
    }

    private void runChains(int length, boolean disableCoerce) {

        Thread[] threads = new Thread[chains.length];
        for (int i = 0; i < chains.length; i++) {
            threads[i] = new MCMCMCRunner(chains[i], length, length, false);
            threads[i].start();
        }

	    // wait for all threads collected to die
	    for (int i =0; i < chains.length; i++) {
		    // wait doggedly for thread to die
		    while (threads[i].isAlive()) {
			    try {
				    threads[i].join();
			    } catch (InterruptedException ie) {
				    // DO NOTHING
			    }
		    }

	    }
    }

    private int swapChainTemperatures() {

        int newColdChain = coldChain;

	    int index1 = MathUtils.nextInt(chains.length);
	    int index2 = MathUtils.nextInt(chains.length);
	    while (index1 == index2) {
		    index2 = MathUtils.nextInt(chains.length);
	    }

	    double score1 = chains[index1].getCurrentScore();
	    MCMCCriterion acceptor1 = ((MCMCCriterion)chains[index1].getAcceptor());
	    double temperature1 = acceptor1.getTemperature();
	    double score2 = chains[index2].getCurrentScore();
	    MCMCCriterion acceptor2 = ((MCMCCriterion)chains[index2].getAcceptor());
	    double temperature2 = acceptor2.getTemperature();

	    double logRatio = ((score2 - score1) * temperature1) + ((score1 - score2) * temperature2);
	    boolean swap = (Math.log(MathUtils.nextDouble()) < logRatio);

	    if (swap) {
		    acceptor1.setTemperature(temperature2);
		    acceptor2.setTemperature(temperature1);

            OperatorSchedule schedule1 = schedules[index1];
            OperatorSchedule schedule2 = schedules[index2];

            for (int i = 0; i < schedule1.getOperatorCount(); i++) {
                MCMCOperator operator1 = schedule1.getOperator(i);
                MCMCOperator operator2 = schedule2.getOperator(i);

                int tmp = operator1.getAccepted();
                operator1.setAccepted(operator2.getAccepted());
                operator2.setAccepted(tmp);

                tmp = operator1.getRejected();
                operator1.setRejected(operator2.getRejected());
                operator2.setRejected(tmp);

                double tmp2 = operator1.getSumDeviation();
                operator1.setDumDeviation(operator2.getSumDeviation());
                operator2.setDumDeviation(tmp2);

                if (operator1 instanceof CoercableMCMCOperator) {
                    tmp2 = ((CoercableMCMCOperator)operator1).getCoercableParameter();
                    ((CoercableMCMCOperator)operator1).setCoercableParameter(((CoercableMCMCOperator)operator2).getCoercableParameter());
                    ((CoercableMCMCOperator)operator2).setCoercableParameter(tmp2);
                }
            }

		    if (index1 == coldChain) {
		        newColdChain = index2;
		    } else if (index2 == coldChain) {
		        newColdChain = index1;
		    }
	    }

        return newColdChain;
    }

    private void resetChains() {

        for (int i = 0; i < chains.length; i++) {
            chains[i].reset();
        }
    }

	/** cleans up when the chain finishes (possibly early). */
	private void finish() {

		NumberFormatter formatter = new NumberFormatter(8);

        MCLogger[] loggers = mcLoggers[coldChain];
        for (int i =0; i < loggers.length; i++) {
            loggers[i].log(currentState);
            loggers[i].stopLogging();
        }

        System.out.println();
        System.out.println("Time taken: " + timer.toString());

	    if (showOperatorAnalysis) {
	        System.out.println();
	        System.out.println("Operator analysis");
	        System.out.println(
	            formatter.formatToFieldWidth("Operator", 30) +
	            formatter.formatToFieldWidth("", 8) +
	            formatter.formatToFieldWidth("Pr(accept)", 11) +
	            " Performance suggestion");
	        for (int i =0; i < schedules[coldChain].getOperatorCount(); i++) {

	            MCMCOperator op = schedules[coldChain].getOperator(i);
	            double acceptanceProb = MCMCOperator.Utils.getAcceptanceProbability(op);
	            String message = "good";
	            if (acceptanceProb < op.getMinimumGoodAcceptanceLevel()) {
	                if (acceptanceProb < (op.getMinimumAcceptanceLevel()/10.0)) {
	                    message = "very low";
	                } else if (acceptanceProb < op.getMinimumAcceptanceLevel()) {
	                    message = "low";
	                } else message = "slightly low";

	            } else if (acceptanceProb > op.getMaximumGoodAcceptanceLevel()) {
	                double reallyHigh = 1.0 - ((1.0-op.getMaximumAcceptanceLevel())/10.0);
	                if (acceptanceProb > reallyHigh) {
	                    message = "very high";
	                } else if (acceptanceProb > op.getMaximumAcceptanceLevel()) {
	                    message = "high";
	                } else message = "slightly high";
	            }

	            String suggestion = op.getPerformanceSuggestion();

	            String pString = "        ";
	            if (op instanceof CoercableMCMCOperator) {
	                pString = formatter.formatToFieldWidth(formatter.formatDecimal(((CoercableMCMCOperator)op).getRawParameter(), 3), 8);
	            }

	            System.out.println(
	                formatter.formatToFieldWidth(op.getOperatorName(), 30) +

	                pString +

	                formatter.formatToFieldWidth(formatter.formatDecimal(acceptanceProb, 4), 11) +
	                 " " + message + "\t" + suggestion);
	        }
	        System.out.println();
	    }
	}

    public class BurninListener implements MarkovChainListener {

        public BurninListener(int stateCount) {
            this.stateCount = stateCount;
            step = 0;
            stepSize = (double)stateCount / 60.0;
        }

        /**
         * Called to update the current model keepEvery states.
         */
        public synchronized void currentState(int state, Model currentModel) {

            if (state == 0) {
                System.out.println();
                System.out.println("Pre-burnin (" + stateCount + " states)");
                System.out.println("0              25             50             75            100");
                System.out.println("|--------------|--------------|--------------|--------------|");
                System.out.print(  "*");
                step = 1;
            }

            if (state >= (int)Math.round(step*stepSize) && step <= 60) {
                System.out.print("*");
                System.out.flush();
                step += 1;
            }
        }

        /** Called when a new new best posterior state is found. */
        public synchronized void bestState(int state, Model bestModel) {}

        /** cleans up when the chain finishes (possibly early). */
        public synchronized void finished(int chainLength) {
            System.out.println("*");
            System.out.println();
        }

        int stateCount = 0;
        double stepSize;
        int step = 0;
    }

    private MarkovChainListener chainListener = new MarkovChainListener() {

        // MarkovChainListener interface *******************************************
        // for receiving messages from subordinate MarkovChain

        /**
         * Called to update the current model keepEvery states.
         */
        public synchronized void currentState(int state, Model currentModel) {

            currentState = state;

            if (state % 1000 == 0) {
                NumberFormatter formatter = new NumberFormatter(8);
                formatter.setPadding(false);

                System.out.print("State " + currentState + ": ");
                for (int i = 0; i < chains.length; i++) {
                    String score;
                    if (i == coldChain) {
                        score = "[" + formatter.format(chains[i].getCurrentScore()) + "]";
                    } else {
                        score = formatter.format(chains[i].getCurrentScore());
                    }
                    score += " ";
                    System.out.print(formatter.formatToFieldWidth(score, 12));
                }
                System.out.println();
            }

            MCLogger[] loggers = mcLoggers[coldChain];
            for (int i =0; i < loggers.length; i++) {
                loggers[i].log(state);
            }

        }

        /** Called when a new new best posterior state is found. */
        public synchronized void bestState(int state, Model bestModel) {
            currentState = state;
        }

        /** cleans up when the chain finishes (possibly early). */
        public synchronized void finished(int chainLength) {
        }

    };

    /** @return the prior of this MCMC analysis. */
    public Prior getPrior() { return chains[coldChain].getPrior(); }

    /** @return the likelihood function. */
    public Likelihood getLikelihood() { return chains[coldChain].getLikelihood(); }

    /** @return the timer. */
    public dr.util.Timer getTimer() { return timer; }

    /** @return the length of this analysis.*/
    public final int getChainLength() { return mcmcOptions.getChainLength(); }

    // TRANSIENT PUBLIC METHODS *****************************************

    /** @return the current state of the MCMC analysis. */
    public final int getCurrentState() { return currentState; }

    /** @return the progress (0 to 1) of the MCMC analysis. */
    public final double getProgress() {
        return (double)currentState / (double)mcmcOptions.getChainLength();
    }

    /** Requests that the MCMC chain stop prematurely. */
    public void pleaseStop() {
        for (int i = 0; i < chains.length; i++) {
            chains[i].pleaseStop();
        }
    }

    //PRIVATE METHODS *****************************************
    private boolean isPreBurninNeeded() {

        if (mcmcOptions.useCoercion()) return true;

        for (int i = 0; i < schedules[coldChain].getOperatorCount(); i++) {
            MCMCOperator op = schedules[coldChain].getOperator(i);

            if (op instanceof CoercableMCMCOperator) {
                if (((CoercableMCMCOperator)op).getMode() == CoercableMCMCOperator.COERCION_ON) return true;
            }
        }
        return false;
    }

    public void setShowOperatorAnalysis(boolean soa) {
        showOperatorAnalysis = soa;
    }

    // PRIVATE TRANSIENTS

	private MCMCOptions mcmcOptions;
	private MCMCMCOptions mcmcmcOptions;

    private boolean showOperatorAnalysis = true;
    private dr.util.Timer timer = new dr.util.Timer();
    private int currentState = 0;

    private MarkovChain[] chains;
    private MCLogger[][] mcLoggers;
    private OperatorSchedule[] schedules;
    private int coldChain;
}

